import os  # 操作系统接口
import joblib  # 模型序列化和保存
import numpy as np  # 数值计算库
import pandas as pd  # 数据处理和分析
from typing import Dict, List, Any, Optional, Tuple  # 类型提示

from sklearn.model_selection import train_test_split, GridSearchCV  # 数据分割和超参数调优
from sklearn.ensemble import RandomForestClassifier  # 随机森林分类器
from sklearn.svm import SVC  # 支持向量机分类器
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score  # 模型评估指标
import matplotlib.pyplot as plt  # 绘图库
import seaborn as sns  # 统计可视化库


class SensorClassifier:
    """
    Class for training and evaluating classification models for sensor data.
    """
    
    def __init__(self, model_type: str = 'rf'):
        """
        Initialize the classifier.
        
        Args:
            model_type: Type of model to use ('rf' for Random Forest, 'svm' for Support Vector Machine)
        """
        self.model_type = model_type  # 模型类型('rf'或'svm')
        self.model = None  # 训练好的模型实例
        self.classes = None  # 类别标签
        
        # 定义GridSearchCV的参数网格
        self.param_grids = {
            'rf': {  # 随机森林参数网格
                'n_estimators': [50, 100, 200],  # 树的数量
                'max_depth': [None, 10, 20, 30],  # 树的最大深度
                'min_samples_split': [2, 5, 10],  # 分裂节点所需最小样本数
                'min_samples_leaf': [1, 2, 4]  # 叶节点所需最小样本数
            },
            'svm': {  # 支持向量机参数网格
                'C': [0.1, 1, 10, 100],  # 正则化参数
                'gamma': ['scale', 'auto', 1, 100],  # 核函数系数
                'kernel': ['rbf', 'linear', 'poly']  # 核函数类型
            }
        }
    
    def train(self, features: pd.DataFrame, labels: List[str]) -> Dict[str, Any]:
        """
        Train the classifier on the given features and labels.
        
        Args:
            features: DataFrame containing features
            labels: List of class labels
            
        Returns:
            Dictionary containing the trained model and evaluation metrics
        """
        # 将数据分割为训练集和测试集(80%训练，20%测试)
        X_train, X_test, y_train, y_test = train_test_split(
            features, labels, test_size=0.2, random_state=42, stratify=labels
        )
        
        # 根据模型类型创建基础模型
        if self.model_type == 'rf':
            base_model = RandomForestClassifier(random_state=42, class_weight='balanced')  # 随机森林
        elif self.model_type == 'svm':
            base_model = SVC(probability=True, random_state=42, class_weight='balanced')  # 支持向量机
        else:
            raise ValueError(f"Unsupported model type: {self.model_type}")
        
        # 使用网格搜索寻找最优超参数
        grid_search = GridSearchCV(
            base_model,
            self.param_grids[self.model_type],  # 使用预定义的参数网格
            cv=5,  # 5折交叉验证
            scoring='accuracy',  # 使用准确率作为评估指标
            verbose=1,  # 显示进度
            n_jobs=-1  # 使用所有CPU核心
        )
        
        grid_search.fit(X_train, y_train)  # 执行网格搜索
        
        # 获取最优模型
        self.model = grid_search.best_estimator_  # 最佳估计器
        self.classes = self.model.classes_  # 类别标签
        
        # 在测试集上进行预测
        y_pred = self.model.predict(X_test)
        
        # 计算评估指标
        accuracy = accuracy_score(y_test, y_pred)  # 准确率
        conf_matrix = confusion_matrix(y_test, y_pred, labels=self.classes)  # 混淆矩阵
        class_report = classification_report(y_test, y_pred, labels=self.classes, output_dict=True)  # 分类报告
        
        # 存储结果
        results = {
            'model': self.model,  # 训练好的模型
            'classes': self.classes,  # 类别标签
            'best_params': grid_search.best_params_,  # 最佳参数
            'accuracy': accuracy,  # 准确率
            'confusion_matrix': conf_matrix,  # 混淆矩阵
            'classification_report': class_report  # 分类报告
        }
        
        return results
    
    def predict(self, features: pd.DataFrame) -> np.ndarray:
        """
        Make predictions on the given features.
        
        Args:
            features: DataFrame containing features
            
        Returns:
            Array of predicted class labels
        """
        if self.model is None:
            raise ValueError("Model has not been trained yet")
        
        return self.model.predict(features)
    
    def predict_proba(self, features: pd.DataFrame) -> np.ndarray:
        """
        Get class probabilities for the given features.
        
        Args:
            features: DataFrame containing features
            
        Returns:
            Array of class probabilities
        """
        if self.model is None:
            raise ValueError("Model has not been trained yet")
        
        return self.model.predict_proba(features)
    
    def save_model(self, model_path: str, results: Dict[str, Any]) -> None:
        """
        Save the trained model and results to a file.
        
        Args:
            model_path: Path to save the model to
            results: Dictionary containing the model and evaluation metrics
        """
        # 如果目录不存在则创建
        os.makedirs(os.path.dirname(model_path), exist_ok=True)
        
        # 使用joblib保存模型和结果
        joblib.dump(results, model_path)  # 序列化并保存整个结果字典
    
    def load_model(self, model_path: str) -> Dict[str, Any]:
        """
        Load a trained model from a file.
        
        Args:
            model_path: Path to the saved model
            
        Returns:
            Dictionary containing the model and evaluation metrics
        """
        results = joblib.load(model_path)  # 从文件加载模型和结果
        
        # 提取模型和类别信息
        self.model = results['model']  # 加载训练好的模型
        self.classes = results['classes']  # 加载类别标签
        
        return results
    
    def plot_confusion_matrix(self, conf_matrix: np.ndarray, classes: List[str],
                             output_path: Optional[str] = None) -> None:
        """
        Plot the confusion matrix.
        
        Args:
            conf_matrix: Confusion matrix
            classes: List of class names
            output_path: Path to save the plot to (if None, the plot will be displayed)
        """
        plt.figure(figsize=(10, 8))  # 设置图形大小
        sns.heatmap(conf_matrix, annot=True, fmt='d', cmap='Blues',  # 创建热力图
                   xticklabels=classes, yticklabels=classes)  # 设置坐标轴标签
        plt.xlabel('Predicted')  # x轴标签(预测值)
        plt.ylabel('True')  # y轴标签(真实值)
        plt.title('Confusion Matrix')  # 图形标题
        
        if output_path:
            # 如果指定了输出路径，则保存图形
            os.makedirs(os.path.dirname(output_path), exist_ok=True)  # 确保目录存在
            plt.savefig(output_path)  # 保存图形到文件
            plt.close()  # 关闭图形释放内存
        else:
            plt.show()  # 否则直接显示图形


if __name__ == "__main__":
    # Test the classifier
    from data_loader import prepare_dataset
    from feature_extraction import extract_features_from_segments, normalize_features
    
    # Load and prepare data
    data_dir = "../data"
    segments, labels, classes = prepare_dataset(data_dir)
    
    # Extract features
    features = extract_features_from_segments(segments)
    
    # Normalize features
    normalized_features, _ = normalize_features(features)
    
    # Train the classifier
    classifier = SensorClassifier(model_type='rf')
    results = classifier.train(normalized_features, labels)
    
    print(f"Accuracy: {results['accuracy']}")
    print(f"Best parameters: {results['best_params']}")
    
    # Plot the confusion matrix
    classifier.plot_confusion_matrix(
        results['confusion_matrix'],
        results['classes'],
        "../output/confusion_matrix.png"
    )
